{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d0a1139b-59a5-4550-90cb-eb1600023f0f",
   "metadata": {},
   "source": [
    "# Demo: HRES T0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d2d6659-d16f-4736-91b0-16efdfa58147",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To run this notebook, you will need a recent version of JupyterLab with iPyLeaflet and CartoPy installed. It is recommended to make a new Conda environment and install all three with `conda install -c conda-forge jupyterlab ipyleaflet cartopy`. In addition, you will need the following PyPI packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74fc66dc-958e-4649-bd5b-4a46c1c9d4a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q fsspec gcsfs \"zarr<3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7320530-4921-4b1f-b0fc-3d6aa6949631",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import warnings\n",
    "\n",
    "# In this demo, we silence all warnings.\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# But we do want to show what's happening under the hood!\n",
    "logging.basicConfig(level=logging.WARNING, format=\"%(asctime)s [%(levelname)s] %(message)s\")\n",
    "logging.getLogger(\"aurora\").setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b4f1375-a591-4208-a45b-4aa47b2864c9",
   "metadata": {},
   "source": [
    "## Load an Initial Condition\n",
    "\n",
    "To run Aurora, you need to initialise the model with the state of the world at a particular point in time. For this demo, we initialise Aurora with HRES T0, which is a state-of-the-art data set that is available in real time. Specifically, we initialise the model with the state of the world at 11 May 2022 UTC 12."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba46076-2b4f-452f-b694-ee108eae839d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from aurora.foundry.demo.hres_t0_data import load_batch\n",
    "\n",
    "initial_condition = load_batch(day=datetime(2022, 5, 11))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20bb6436-b26e-4aca-a354-4dfbc60b37a8",
   "metadata": {},
   "source": [
    "## Produce Prediction on Azure AI Foundry\n",
    "\n",
    "Now that we have an initial condition, we connect to the Foundry endpoint and run the model for 4 6-hour steps, producing predictions up to and including UTC 12 the next day.\n",
    "\n",
    "To run on Foundry, the environment variables `FOUNDRY_ENDPOINT`, `FOUNDRY_TOKEN`, and `BLOB_URL_WITH_SAS` need to be set. If you're unsure on how to set environment variables, see [here](intro.md), or simply copy paste your keys below (but be sure not to commit any keys)! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07fd311b-54c0-4144-beb6-4a81d2056507",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from aurora.foundry import BlobStorageChannel, FoundryClient, submit\n",
    "\n",
    "foundry_client = FoundryClient(\n",
    "    endpoint=os.environ[\"FOUNDRY_ENDPOINT\"],\n",
    "    token=os.environ[\"FOUNDRY_TOKEN\"],\n",
    ")\n",
    "channel = BlobStorageChannel(os.environ[\"BLOB_URL_WITH_SAS\"])\n",
    "\n",
    "predictions = list(\n",
    "    submit(\n",
    "        initial_condition,\n",
    "        model_name=\"aurora-0.25-finetuned\",\n",
    "        num_steps=4,  # Every step predicts 6 hours into the future.\n",
    "        foundry_client=foundry_client,\n",
    "        channel=channel,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb37746-d2d9-4715-8ddf-37252a5a3190",
   "metadata": {},
   "source": [
    "## Visualisation of Predictions\n",
    "\n",
    "These are the predictions for temperature globally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ea3d81c-6932-46f6-a91b-6fc79d549b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def plot(ax, data, time, ic=True):\n",
    "    ax.imshow(data.numpy() - 273.15, vmin=-50, vmax=50, cmap=\"RdBu_r\")\n",
    "    ax.set_title(f\"2T at {time} \" + (\"(initial condition)\" if ic else \"(prediction)\"))\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(3, 2, figsize=(10, 8))\n",
    "ic = initial_condition\n",
    "plot(axs[0, 0], ic.surf_vars[\"2t\"][0, 0], ic.metadata.time[0] - timedelta(hours=6))\n",
    "plot(axs[0, 1], ic.surf_vars[\"2t\"][0, 1], ic.metadata.time[0])\n",
    "for i in range(2, 6):\n",
    "    pred = predictions[i - 2]\n",
    "    plot(axs[i // 2, i % 2], pred.surf_vars[\"2t\"][0, 0], pred.metadata.time[0], ic=False)\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(hspace=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f611b7fc-a4d8-4fe6-9650-b5cf2a817726",
   "metadata": {},
   "source": [
    "Aurora, however, predicts much more than just temperature: the model predicts the entire state of the atmosphere! Here is an interactive plot of the predictions for wind speed, temperature, and the mean pressure at sea level in the first step of the prediction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876d3ec7-7dfe-45ac-852d-96aac1f0ef46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aurora.foundry.demo.interactive import interactive_plot\n",
    "\n",
    "interactive_plot(predictions[0], width=\"1000px\", height=\"500px\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daaefadb-308b-4a24-8635-9d00cc78aaf7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aurora-dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
